package org.znxs.znagent_s.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.UUID;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import org.znxs.znagent_s.agent.model.ZNAgentStatus;

import java.util.List;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * 工具调用实体类
 * 智能体工具调用 实体类 负责实现工具调用能力
 * <p>
 * 因为Spring ai 封装的内部的判断是否需要调用工具，
 * 如果我们需要自己实现手动调用工具，就需要采用另一种方式
 * <p>
 * 1）基于Spring AI 的工具调用能力，手动控制工具执行。
 * <p>
 * 其实 Spring 的 ChatClient 已经支持选择工具进行调用（内部完成了 think、act、observe），
 * 但这里我们要自主实现，可以使用 Spring AI 提供的 手动控制工具执行。
 * <p>
 * 但是Spring AI 这方面有些问题，因为我们采用的是阿里巴巴的Spring ai 所以使用spring 这方面的功能会有兼容性问题
 * 我们通过将 DashScopeChatOptions 的 withProxyToolCalls 选项设置为 true，来禁止 Spring AI 托管工具调用，而是我们自主实现。
 */
@Slf4j
public class ZNToolCallAgent extends ZNReActAgent {

    /**
     * 可调用的方法集合
     */
    private final ToolCallback[] availableTools;


    /**
     * 保存了工具调用信息的响应
     */
    private ChatResponse toolCallChatResponse;

    /**
     * 工具调用管理者
     */
    private final ToolCallingManager toolCallingManager;

    /**
     * 对话选项 通过禁用内置的工具调用机制，自己维护上下文
     */
    private final ChatOptions chatOptions;


    public ZNToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        // 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true) // 设置自己代理工具调用 从而不适用Spring AI 的工具调用
                .build();

        // 使用下面的官方代码会直接导致工具调用失效！
        /*ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(new CustomerTools())
                .internalToolExecutionEnabled(false)
                .build();*/
    }


    /**
     * 处理当前状态并决定下一步行动
     *
     * @return 是否需要执行行动
     */
    @Override
    public Boolean think() {
        // 下一步的单步预设prompt
        if (getNextPrompt() != null && !getNextPrompt().isEmpty()) {
            UserMessage userMessage = new UserMessage(getNextPrompt());
            getMessages().add(userMessage);
        }
        // 获取消息列表
        List<Message> messageList = getMessages();
        // 封装提示实体 将chatOptions封装进去，实现自己控制调用工具的选择
        Prompt prompt = new Prompt(messageList, chatOptions);
        // 向ai进行提问 判断是否使用工具
        try {
            //对话记忆的唯一标识
            String conversantId = UUID.randomUUID().toString();
            // 获取带工具选项的响应
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt()) // 系统预设prompt
                    .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversantId)
                            .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                    .tools(availableTools) // 工具调用
                    .call()
                    .chatResponse();
            // 记录响应，用于 Act
            this.toolCallChatResponse = chatResponse;
            // 提取出返回结果信息
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            // 输出提示信息
            String result = assistantMessage.getText();
            // 调用工具的集合
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
            log.info(getAgentName() + "的思考：" + result);
            log.info(getAgentName() + "选择了：" + toolCalls.size() + " 个工具来调用");
            String toolCallInfo = toolCalls.stream().map(toolCall ->
                    String.format("工具名称：%s，参数%s", toolCall.name(), toolCall.arguments())).collect(Collectors.joining("\n"));
            log.info(toolCallInfo);

            //只有不调用工具的时候才记录助手消息
            if (toolCalls.isEmpty()) {
                getMessages().add(assistantMessage);
                return false;
            } else {
                // 需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        } catch (Exception e) {
            // 遇到问题
            log.error(getAgentName() + "思考时遇到问题：", e.getMessage());
            getMessages().add(new AssistantMessage("思考处理时遇到问题" + e.getMessage()));
            return false;
        }

    }


    /**
     * 执行工具调用并处理结果
     *
     * @return 执行结果
     */
    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "没有工具调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(getMessages(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        // 记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        setMessages(toolExecutionResult.conversationHistory());
        // 当前工具调用结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        // 提取调用结果作为返回结果
        String results = toolResponseMessage.getResponses().stream()
                .map(toolResponse ->
                        "工具：" + toolResponse.name() + "，完成了他的任务！结果为：" + toolResponse.responseData())
                .collect(Collectors.joining("\n"));

        // 判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> "doTerminate".equals(response.name()));
        if (terminateToolCalled) {
            setStatus(ZNAgentStatus.FINISHED);
        }
        log.info(results);
        return results;
    }
}
